package com.llmops.demo;

import com.llmops.core.Context;
import com.llmops.core.Edge;
import com.llmops.core.WorkflowGraph;
import com.llmops.core.node.ForNode;
import com.llmops.core.node.StartNode;

import java.util.Collections;
import java.util.List;
import java.util.Map;

public class ForNodeTest {
    public static void main(String[] args) {
        WorkflowGraph workflow = new WorkflowGraph("1");
        Context context = new Context();
        context.setVariable("sessionId",3);
        context.setVariable("list_input", List.of("AA","BB","CC"));
        workflow.addNode(new StartNode("data_input",Collections.singletonMap("input","你是谁")));


        WorkflowGraph forWorkflow = new WorkflowGraph("forWorkflow");
        forWorkflow.addNode(new StartNode("forWorkflow_data_input",null));

//        workflow.addNode(new LLMNode("llm_inferenceB", Collections.singletonMap("model", "gpt-B")));

        workflow.addNode(new ForNode("forNode", Map.of("maxTimes",2,"listName","list_input"),forWorkflow));

        workflow.addEdge(new Edge("data_input","forNode"));
        workflow.validate();
//        workflow.execute();

        workflow.execute(context);

        System.out.println("Workflow execution completed");
    }
}